#!/bin/bash -e
#
# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

tests_dir=${1:-"@PYTHON_TEST_DIR@"}
python_versions_tentative="@PYTHON_TEST_VERSIONS@"

python_versions=""

# Check if test scripts exist in test_dir
if [ ! -f "$tests_dir/cvcuda_util.py" ]; then
    # try to find the scripts in relative path
    echo "Python test scripts (E.g. cvcuda_util.py, test_op*.py) doesn't exist in $tests_dir"
    echo "Trying to find python test scripts via relative path"
    cvcuda_test_types_python_tar_dir=$(cd "$(dirname "$0")"; pwd)/../@PYTHON_TEST_INSTDIR@ # relative path in tarball
    if [ -f "$cvcuda_test_types_python_tar_dir/cvcuda_util.py" ]; then
        echo "Found python test scripts at $cvcuda_test_types_python_tar_dir via relative path"
        tests_dir=$cvcuda_test_types_python_tar_dir
    else
        echo "Cannot find python test scripts in $tests_dir and $cvcuda_test_types_python_tar_dir"
        echo "Please run ./cvcuda_test_python [python test scripts folder]"
        exit 1 #hard exit
    fi
fi

# Verify if correct package dependencies are installed --------
pip_depends="pytest torch numpy==1.26"

# Collect all python versions that are indeed installed and have proper dependencies installed
# Two behaviors:
# - default: skip Python versions that are not installed or don't have pytest and torch installed
# - if NVCV_FORCE_PYTHON is set: exit with error
for ver in $python_versions_tentative; do
    if ! python$ver -c "import pytest, torch, numpy" > /dev/null 2>&1; then
        echo "WARNING: Python version $ver not installed or missing proper dependencies"
        echo "Please install Python version $ver and run the following commands before running tests: python$ver -m pip install $pip_depends"
        if [[ "$NVCV_FORCE_PYTHON" == 1 || "$NVCV_FORCE_PYTHON" == yes ]]; then
            echo "Exiting with FAILURE, as NVCV_FORCE_PYTHON=$NVCV_FORCE_PYTHON"
            exit 2 #hard exit
        else
            echo "Continue and skipping python version $ver, as NVCV_FORCE_PYTHON=$NVCV_FORCE_PYTHON"
        fi
    else
        echo "Found Python version $ver installed with proper dependencies, adding to tests"
        python_versions+="$ver "
    fi
done

# Run tests --------

tmpdir=$(mktemp -d)

function on_exit()
{
    rm -rf $tmpdir
}
trap 'on_exit' EXIT

export PYTHONPATH="$PYTHONPATH:@PYTHON_MODULE_DIR@"

for ver in $python_versions; do

    if [[ "$NVCV_FORCE_PYTHON" != 1 && "$NVCV_FORCE_PYTHON" != yes ]]; then
        if ! python$ver -c 'import cvcuda'; then
            echo "Skipping python-$ver, CV-CUDA python bindings not installed"
            continue
        fi
    fi

    # Check if python module is exposing only PyInit_cvcuda.
    # Also provide some helpful info is exposing too much.
    modfile=$(python$ver -c "import cvcuda; print(cvcuda.__file__)")
    pubsyms=$(readelf -sWD $modfile | grep -v ' UND ' | grep ' GLOBAL ')
    if [[ $(echo "$pubsyms" | wc -l) != 1 ]]; then
        echo -e "cvcuda python $ver module is exposing too many symbols:\n$pubsyms"
        exit 3
    fi
    if ! echo "$pubsyms" | grep PyInit_cvcuda > /dev/null; then
        echo -e "cvcuda python $ver module must expose symbol PyInit_cvcuda, but instead exposes:\n$pubsyms"
        exit 4
    fi

    # Run python tests
    NVCV_VERSION="@NVCV_VERSION_FULL@" python$ver -m pytest -v --tb=line -o cache_dir="$tmpdir" "$@" "$tests_dir"
done
